import pandas as pd
import matplotlib.pyplot as plt
import matplotlib as mpl

# 设置全局字体
plt.rcParams['font.family'] = 'Times New Roman'  # 默认使用Times New Roman
plt.rcParams['axes.unicode_minus'] = False

# 创建字体属性对象用于中文显示
chinese_font = {'family': 'SimHei', 'size': 14}  # 中文使用黑体

# 读取数据
df = pd.read_csv("mtcars1.csv")

# 计算标准化的 mpg 值
x = df.loc[:, ['mpg']]
df['mpg_z'] = (x - x.mean()) / x.std()

# 定义颜色映射函数
def assign_color(z):
    if z < 0:
        return '#f0b2ab'  # 自定义粉色（浅粉）
    else:
        return '#95c8e7'  # 自定义蓝色（浅蓝）

# 应用颜色函数
df['colors'] = df['mpg_z'].apply(assign_color)

# 按 mpg_z 排序，重置索引
df.sort_values('mpg_z', inplace=True)
df.reset_index(inplace=True)

# 绘图
plt.figure(figsize=(10, 8), dpi=80)
plt.hlines(y=df.index, xmin=0, xmax=df.mpg_z, color=df.colors, alpha=0.7, linewidth=8)

# 图形美化设置
plt.gca().set(ylabel='Model', xlabel='Mileage')
plt.yticks(df.index, df['cars'], fontsize=14)
plt.title('xxx蝴蝶图', fontdict={'fontproperties': chinese_font, 'size': 20})
plt.xlabel('里程数（单位：万）', fontdict={'fontproperties': chinese_font, 'size': 14})
plt.ylabel('车型', fontdict={'fontproperties': chinese_font, 'size': 14})

# 设置坐标轴标签字体为Times New Roman
for label in plt.gca().get_xticklabels() + plt.gca().get_yticklabels():
    label.set_fontname('Times New Roman')
    label.set_fontsize(14)

# 在条形上添加标签
for x, y, tex in zip(df.mpg_z, df.index, df.mpg_z):
    t = plt.text(x, y, round(tex, 2),
                 horizontalalignment='right' if x < 0 else 'left',
                 verticalalignment='center',
                 fontdict={'color': '#f0b2ab' if x < 0 else '#95c8e7', 'size': 12})

#设置网格线
plt.grid(linestyle='--', alpha=0.5)
plt.xlim(-2.3,2.3) #防止标签和坐标轴重合，范围设置大了一点点

# 保存图像
plt.savefig('蝴蝶图.png', dpi=300, bbox_inches='tight')
plt.show()